# -*- coding:utf-8 -*-

'''
python-ldap包的用法
use it to access lenovo AD LDAP

pip install python-ldap

ldap.contextSource.url=ldap://lenovo.com:389
ldap.contextSource.base=OU=User Accounts,DC=lenovo,DC=com
ldap.contextSource.userDn=不告诉你
ldap.contextSource.password=不告诉你
ldap.contextSource.query=cn

CN=ctoprd01, OU=Service Accounts, DC=lenovo, DC=com



'''

import ldap
import json


class Employee(object):

    def __init__(self, dn, name):
        self.dn = dn
        self.name = name
        self.members = None

    def get_members(self, lenovo_ad):
        if self.dn:
            mems = lenovo_ad.search_members_of_manager(self.dn)
            self.members = []
            for m in mems:
                e = Employee(m[0], m[1]['cn'])
                self.members.append(e)
                e.get_members(lenovo_ad)

    def json_encoder_default(self, o):
        return {'dn': self.dn, 'name': self.name, 'members': []}


class LenovoAD(object):

    def __init__(self, serverUrl, base_dn, user_dn, password, attrs=None):
        self.serverUrl = serverUrl
        self.base_dn = base_dn
        self.user_dn = user_dn
        self.password = password
        self.ldap_obj = None
        if attrs:
            self.attrs = attrs
        else:
            self.attrs = ['cn', 'mail', 'distinguishedName', 'telephoneNumber',
                          'department', 'manager']

    def init(self):
        '''初始化ldap对象，初始化后就可以做查询操作
        '''
        self.ldap_obj = ldap.initialize(self.serverUrl)
        self.ldap_obj.set_option(ldap.OPT_REFERRALS, 0)
        self.ldap_obj.protocol_version = ldap.VERSION3
        self.ldap_obj.simple_bind_s(self.user_dn, self.password)

    def search_by_itcode(self, itcode, attrs=None):
        '''根据itcode搜索组织架构
        for example: cn_name is cn=zhaoxp2
        return list of objects maching cn_name
        '''
        print('search by itcode %s' % itcode)
        if attrs:
            search_attrs = attrs
        else:
            search_attrs = self.attrs
        res = self.ldap_obj.search_s(self.base_dn,
                                     ldap.SCOPE_SUBTREE,
                                     'cn=%s' % itcode,
                                     search_attrs)
        return self._process_search_result(res)

    def search_members_of_manager(self, manager_dn, attrs=None):
        '''搜索孩子，返回dict的列表，其中行key是attrs中定义的值，value是对应的值
        '''
        # print('search members of %s ' % manager_dn)
        if attrs:
            search_attrs = attrs
        else:
            search_attrs = self.attrs
        res = self.ldap_obj.search_s(self.base_dn,
                                     ldap.SCOPE_SUBTREE,
                                     '(&(manager=%s)(employeeType=Regular '
                                     'Employee))' % manager_dn,
                                     search_attrs)
        return self._process_search_result(res)
    
    def _process_search_result(self, res):
        '''处理LDAP搜索结果，将结果转换为dict的list，其中dict记录了dn/cn等信息
        '''
        children = []
        for item in res:
            child = {'dn': item[0]}
            kvs = item[1]
            for key in kvs:
                child[key] = str(kvs[key][0], encoding='utf-8')
            children.append(child)
        return children

    def generate_org_tree(self, itcode, attrs=None):
        '''生产组织关系的树形结构数据，树形结构用dict表示。
        其中孩子对应的key为members，值为list。
        '''
        # print('generate org tree for %s' % itcode)
        res = self.search_by_itcode(itcode, attrs)
        if res:
            if len(res) == 1:
                return self._get_members(res[0], attrs)
            else:
                raise ValueError('more thant 1 result for %s' % itcode)
        return None
    
    def _get_members(self, user_info, attrs=None):
        '''获得user_info对应用户的下属members。
        其中，user_info是一个dict，里面必须包含了dn和cn。
        attrs代表了查询哪些属性。
        attrs为空，有孩子：返回dict，key为itcode，value值为孩子。
        attrs为空，无孩子：返回itcode。
        attrs不为空，有孩子：返回dict，包含了attrs属性和包含孩子的member属性。
        attrs不为空，无孩子：返回dict，包含了attrs属性。
        '''
        # print('  ------> %s' % user_info)
        user_dn = user_info['dn']
        it_code = user_info['cn']
        members = self.search_members_of_manager(user_dn, attrs)
        # print('%s member size is %s' % (itcode, len(members)))
        # chart = '  %s' % it_code
        if members:
            mems = []
            for m in members:
                # chart += '\n    |-%s' % m['cn']
                mems.append(self._get_members(m, attrs))
            # print(chart)
            # return {'itcode': itcode, 'dn': user_dn, 'members': mems}
            if attrs:
                user_info.update({'members': mems})
                return user_info
            else:
                return {it_code: mems}
        else:
            # print(chart)
            # return {'itcode': itcode, 'dn': user_dn}
            if attrs:
                return user_info
            else:
                return it_code


def main():
    print('test ldap')
    server = 'ldap://lenovo.com:389'
    base_dn = 'OU=User Accounts,DC=lenovo,DC=com'
    user_dn = 'cn=不告诉你,OU=Service Accounts,DC=lenovo,DC=com'
    password = '不告诉你'
    attrs = ['cn', 'mail', 'distinguishedName', 'telephoneNumber', 'department', 'manager']
    lenovo_ad = LenovoAD(server, base_dn, user_dn, password)
    lenovo_ad.init()
    username = 'yaozy'
    res = lenovo_ad.search_by_itcode(username, attrs)
    if res:
        print('%s found -> %s' % (username, res))
        result = lenovo_ad.generate_org_tree(username, attrs)
        print('member tree is ------>')
        print(json.dumps(result, indent=4))
    else:
        print('%s not found in LenovoAD' % username)


if __name__ == '__main__':
    main()
